#pragma once

using namespace System::Net;
using namespace System::Net::Sockets;

namespace ManagedWin32
{
	static_assert((int)AddressFamily::InterNetwork == AF_INET, "");
	static_assert((int)AddressFamily::InterNetworkV6 == AF_INET6, "");
	static_assert((int)SocketType::Stream == SOCK_STREAM, "");
	static_assert((int)SocketType::Dgram == SOCK_DGRAM, "");
	static_assert((int)ProtocolType::Tcp == IPPROTO_TCP, "");
	static_assert((int)ProtocolType::Udp == IPPROTO_UDP, "");
	static_assert((int)SocketShutdown::Send == SD_SEND, "");

	union SocketAddress
	{
		struct sockaddr sa;
		struct sockaddr_in sa_in;
		struct sockaddr_in6 sa_in6;
	};

	SocketAddress CreateSocketAddressFromIPEndPoint(IPEndPoint^ endPoint)
	{
		SocketAddress socketAddress;
		memset(&socketAddress, 0, sizeof(SocketAddress));
		IPAddress^ address = endPoint->Address;
		array<System::Byte>^ addressBytes = address->GetAddressBytes();
		pin_ptr<System::Byte> pinnedAddressBytes = &addressBytes[0];
		switch (endPoint->AddressFamily)
		{
		case AddressFamily::InterNetwork:
			socketAddress.sa_in.sin_family = AF_INET; 
			socketAddress.sa_in.sin_port = htons(endPoint->Port);
			memcpy(&socketAddress.sa_in.sin_addr, pinnedAddressBytes, sizeof(socketAddress.sa_in.sin_addr));
			break;
		case AddressFamily::InterNetworkV6:
			socketAddress.sa_in6.sin6_family = AF_INET6; 
			socketAddress.sa_in6.sin6_port = htons(endPoint->Port);
			memcpy(&socketAddress.sa_in6.sin6_addr, pinnedAddressBytes, sizeof(socketAddress.sa_in6.sin6_addr));
			break;
		default:
			throw gcnew ArgumentException("Unsupported AddressFamily: " + endPoint->AddressFamily.ToString());
		}
		return socketAddress;
	}
	
	int GetSocketAddressLength(const SocketAddress &socketAddress)
	{
		switch (socketAddress.sa.sa_family) 
		{
		case AF_INET:
			return sizeof(socketAddress.sa_in);
		case AF_INET6:
			return sizeof(socketAddress.sa_in6);
		default:
			return sizeof(socketAddress.sa);
		}
	}

	IPEndPoint^ CreateIPEndPointFromSocketAddress(const SocketAddress &socketAddress)
	{
		const void *addressPtr = nullptr;
		int addressLength = 0;
		int port = 0;
		switch (socketAddress.sa.sa_family)
		{
		case AF_INET:
			addressPtr = &socketAddress.sa_in.sin_addr;
			addressLength = sizeof(socketAddress.sa_in.sin_addr);
			port = htons(socketAddress.sa_in.sin_port);
			break;
		case AF_INET6:
			addressPtr = &socketAddress.sa_in6.sin6_addr;
			addressLength = sizeof(socketAddress.sa_in6.sin6_addr);
			port = htons(socketAddress.sa_in6.sin6_port);
			break;
		}
		array<System::Byte>^ addressBytes = gcnew array<System::Byte>(addressLength);
		pin_ptr<System::Byte> pinnedAddressBytes = &addressBytes[0];
		memcpy(pinnedAddressBytes, addressPtr, addressLength);
		IPAddress^ address = gcnew IPAddress(addressBytes);
		IPEndPoint^ endPoint = gcnew IPEndPoint(address, port);
		return endPoint;
	}

	public ref class Win32Socket : public IDisposable
	{
	public:

		Win32Socket(AddressFamily addressFamily, SocketType socketType, ProtocolType protocolType)
		{
			m_handle = WSASocketA((int)addressFamily, (int)socketType, (int)protocolType, NULL, 0, 0);
			if (m_handle == INVALID_SOCKET) {
				throw gcnew SocketException();
			}
		}

		~Win32Socket()
		{
			this->!Win32Socket();
		}

		!Win32Socket()
		{
			if (m_handle != INVALID_SOCKET) {
				closesocket(m_handle);
				m_handle = INVALID_SOCKET;
			}
		}

		void Bind(EndPoint^ endPoint)
		{
			if (endPoint == nullptr) {
				throw gcnew ArgumentNullException();
			}
			SocketAddress socketAddress = CreateSocketAddressFromIPEndPoint((IPEndPoint^) endPoint);
			if (bind(m_handle, &socketAddress.sa, GetSocketAddressLength(socketAddress)) != 0) {
				throw gcnew SocketException();
			}
		}

		void Listen(int backlog)
		{
			if (listen(m_handle, backlog) != 0) {
				throw gcnew SocketException();
			}
		}

		Win32Socket^ Accept() 
		{
			SOCKET handle = accept(m_handle, NULL, NULL);
			if (handle == INVALID_SOCKET) {
				throw gcnew SocketException();
			}
			return gcnew Win32Socket(handle);
		}

		void Connect(EndPoint^ endPoint) 
		{
			if (endPoint == nullptr) {
				throw gcnew ArgumentNullException();
			}
			SocketAddress socketAddress = CreateSocketAddressFromIPEndPoint((IPEndPoint^) endPoint);
			if (connect(m_handle, &socketAddress.sa, GetSocketAddressLength(socketAddress)) != 0) {
				throw gcnew SocketException();
			}
		}

		void Shutdown(SocketShutdown how) 
		{
			if (shutdown(m_handle, (int)how) != 0) {
				throw gcnew SocketException();
			}
		}

#define PRE_SEND_RECEIVE() \
		if (buffer == nullptr) { \
			throw gcnew ArgumentNullException(); \
		} \
		if ((offset < 0) || (count < 0) || (offset + count < 0) || (offset + count > buffer->Length)) { \
			throw gcnew ArgumentOutOfRangeException(); \
		} \
		pin_ptr<System::Byte> data = &buffer[offset]; \
		
		int Send(array<System::Byte>^ buffer, int offset, int count)
		{
			PRE_SEND_RECEIVE();
			int result = send(m_handle, (const char *) data, count, 0);
			if (result == SOCKET_ERROR) {
				throw gcnew SocketException();
			}
			return result;
		}

		int Send(array<System::Byte>^ buffer)
		{
			return Send(buffer, 0, buffer->Length);
		}

		int Receive(array<System::Byte>^ buffer, int offset, int count)
		{
			PRE_SEND_RECEIVE();
			int result = recv(m_handle, (char *) data, count, 0);
			if (result == SOCKET_ERROR) {
				throw gcnew SocketException();
			}
			return result;
		}

		int Receive(array<System::Byte>^ buffer)
		{
			return Receive(buffer, 0, buffer->Length);
		}

		int SendTo(array<System::Byte>^ buffer, int offset, int count, EndPoint^ endPoint)
		{
			if (endPoint == nullptr) {
				throw gcnew ArgumentNullException();
			}
			PRE_SEND_RECEIVE();
			SocketAddress socketAddress = CreateSocketAddressFromIPEndPoint((IPEndPoint^) endPoint);
			int result = sendto(m_handle, (const char *) data, count, 0, 
				&socketAddress.sa, GetSocketAddressLength(socketAddress));
			if (result == SOCKET_ERROR) {
				throw gcnew SocketException();
			}
			return result;
		}

		int SendTo(array<System::Byte>^ buffer, EndPoint^ endPoint)
		{
			return SendTo(buffer, 0, buffer->Length, endPoint);
		}

		int ReceiveFrom(array<System::Byte>^ buffer, int offset, int count, EndPoint^% endPoint)
		{
			PRE_SEND_RECEIVE();
			SocketAddress socketAddress;
			int socketAddressLength = sizeof(SocketAddress);
			int result = recvfrom(m_handle, (char *) data, count, 0,
				&socketAddress.sa, &socketAddressLength);
			if (result == SOCKET_ERROR) {
				throw gcnew SocketException();
			}
			endPoint = CreateIPEndPointFromSocketAddress(socketAddress);
			return result;
		}

		int ReceiveFrom(array<System::Byte>^ buffer, EndPoint^% endPoint)
		{
			return ReceiveFrom(buffer, 0, buffer->Length, endPoint);
		}

	private:
		
		SOCKET m_handle;

		Win32Socket(SOCKET handle)
			: m_handle(handle)
		{
		}
	};

	public ref class Win32SocketStream : Stream
	{
	private:

		Win32Socket^ m_socket;

	public:

		Win32SocketStream(Win32Socket^ socket)
		{
			if (socket == nullptr) {
				throw gcnew ArgumentNullException();
			}
			m_socket = socket;
		}

		property bool CanRead
		{
			virtual bool get() override { return true; }
		}

		property bool CanSeek
		{
			virtual bool get() override { return false; }
		}

		property bool CanWrite
		{
			virtual bool get() override { return true; }
		}

		property __int64 Length
		{
			virtual __int64 get() override { throw gcnew NotSupportedException(); }
		}

		property __int64 Position
		{
			virtual __int64 get() override { throw gcnew NotSupportedException(); }
			virtual void set(__int64) override { throw gcnew NotSupportedException(); }
		}

		virtual void Flush() override
		{
			// Nothing to do.
		}

		virtual __int64 Seek(__int64 offset, SeekOrigin origin) override
		{
			throw gcnew NotSupportedException();
		}

		virtual void SetLength(__int64 value) override
		{
			throw gcnew NotSupportedException();
		}

		virtual int Read(array<System::Byte>^ bytes, int offset, int count) override
		{
			return m_socket->Receive(bytes, offset, count);
		}

		virtual void Write(array<System::Byte>^ bytes, int offset, int count) override
		{
			if (m_socket->Send(bytes, offset, count) != count) {
				throw gcnew IOException("Incomplete socket send operation");
			}
		}
	};
}
